#!/usr/bin/env python

import json
from entry import create_payload, create_inputs
import sys
import zipfile
import io
import platform
import urllib2


def usage():
	print("""Usage: %s <command>

Available commands:
  <payload_file>                       upload a file and run
  --overwrite <target> [<src>]         upload a file to target
  --transmission-add-torrent <files..> upload torrent files
  --transmission-add-url <urls...>     add urls, same as --transmission -a
  --transmission-list-torrent          list torrents and files
  --transmission <commands...>         run transmission commands, for example
    -a, --add <filenames-or-URLs>      add torrents to transmission
    -r, --remove                       remove torrents while keeping data
    --remove-and-delete                remove torrents and delete data
    -s, --start                        start torrents
    -S, --stop                         stop torrents
    -t, --torrent all | <id>           select torrents for subsequent options
  --upload-by-scp <remote_files...>    upload a remote file or dir by scp""" % sys.argv[0])
	sys.exit(1)


def main():
	if len(sys.argv) < 2:
		usage()

	if sys.argv[1].startswith("--"):
		if sys.argv[1] == "--transmission-list-torrent":
			print("* Operation: list torrents and files")
			payload = '#!/bin/bash\n'
			payload += 'df\n'
			payload += 'transmission-remote "$TRANSMISSION_SERVER" -n "$TRANSMISSION_USERPASS" -l\n'
			payload += 'ls -l -R -L "$TRANSMISSION_DIRECTORY"\n'
			ext = 'sh'
		elif sys.argv[1] == "--transmission":
			if len(sys.argv) < 3:
				usage()
			# FIXME: the commands can't contain '"', '\', etc.
			cmd = '"' + '" "'.join(sys.argv[2:]) + '"'
			print("* Operation: run transmission command: " + cmd)
			payload = '#!/bin/bash\n'
			payload += 'df\n'
			payload += 'transmission-remote "$TRANSMISSION_SERVER" -n "$TRANSMISSION_USERPASS" ' + cmd + '\n'
			payload += 'ls -l -R -L "$TRANSMISSION_DIRECTORY"\n'
			ext = 'sh'
		elif sys.argv[1] == "--transmission-add-url":
			if len(sys.argv) < 3:
				usage()
			payload = '#!/bin/bash\n'
			for s in sys.argv[2:]:
				print("* Operation: add torrent url %s" % s)
				payload += 'transmission-remote "$TRANSMISSION_SERVER" -n "$TRANSMISSION_USERPASS" -a "%s"\n' % s
			payload += 'df\n'
			payload += 'transmission-remote "$TRANSMISSION_SERVER" -n "$TRANSMISSION_USERPASS" -l\n'
			payload += 'ls -l -R -L "$TRANSMISSION_DIRECTORY"\n'
			ext = 'sh'
		elif sys.argv[1] == "--upload-by-scp":
			if len(sys.argv) < 3:
				usage()
			payload = '#!/bin/bash\n'
			for tf in sys.argv[2:]:
				target_file = tf.replace("\\", "/")
				target_file = target_file.decode(sys.getfilesystemencoding()).encode('utf-8') # assume target system is UTF-8
				try:
					print("* Operation: upload %s by scp\n" % target_file)
				except:
					pass
				# FIXME: if filename contains '`' then it doesn't work
				target_file = target_file.replace("`", "\\`") # ???
				if '?' in target_file or '*' in target_file:
					lps = target_file.rfind('/')
					payload += 'cd "%s"\n' % target_file[:lps+1]
					payload += 'find . -maxdepth 1 -name "%s" -exec echo "* Uploading" {} "..." \\; -exec time scp $SCP_CMD_OPTIONS -r {} "$SCP_TARGET" \\;\n' % target_file[lps+1:]
				else:
					payload += 'echo "* Uploading %s ..."\n' % target_file
					payload += 'time scp $SCP_CMD_OPTIONS -r "%s" "$SCP_TARGET"\n' % target_file
			ext = 'sh'
		else:
			with io.BytesIO() as fp:
				with zipfile.ZipFile(fp, 'w', zipfile.ZIP_STORED, False) as zf:
					if sys.argv[1] == "--overwrite":
						if len(sys.argv) < 3:
							usage()
						target_file = sys.argv[2]
						if len(sys.argv) < 4:
							src_file = target_file
						else:
							src_file = sys.argv[3]
						if platform.system() == "Windows":
							src_file = src_file.replace("/", "\\")
						else:
							src_file = src_file.replace("\\", "/")
						target_file = target_file.replace("\\", "/")
						target_file = target_file.decode(sys.getfilesystemencoding()).encode('utf-8') # assume target system is UTF-8
						try:
							print("* Operation: upload file %s to %s" % (src_file, target_file))
						except:
							pass
						zf.write(src_file, "temp.bin")
						payload = '#!/bin/bash\n'
						payload += 'mv -f build/payload/temp.bin "%s"\n' % target_file
						zf.writestr("payload.sh", payload)
					elif sys.argv[1] == "--transmission-add-torrent":
						if len(sys.argv) < 3:
							usage()
						filecount = 0
						for fn in sys.argv[2:]:
							print("* Operation: upload torrent file %s" % fn)
							filename_lower = fn.lower()
							if filename_lower.startswith("http://") or filename_lower.startswith("https://"):
								# download file first
								print("Downloading torrent file...")
								try:
									data = urllib2.urlopen(fn).read()
									# check file format
									if data[:11] != "d8:announce":
										print("ERROR: file is not a valid torrent file")
									else:
										zf.writestr("temp%03d.torrent" % filecount, data)
										filecount += 1
								except Exception as e:
									print "ERROR: Failed to download torrent file: ", e
							else:
								try:
									# check file format
									with open(fn, 'rb') as ff:
										torrent_file_header = ff.read(11)
									if torrent_file_header != "d8:announce":
										print("ERROR: file is not a valid torrent file")
									else:
										zf.write(fn, "temp%03d.torrent" % filecount)
										filecount += 1
								except Exception as e:
									print "ERROR: Failed to open torrent file: ", e
						print("Found %d torrent file(s)" % filecount)
						if filecount == 0:
							print("ERROR: No files to upload, quitting")
							sys.exit(-1)
						payload = '#!/bin/bash\n'
						for i in range(filecount):
							payload += 'transmission-remote "$TRANSMISSION_SERVER" -n "$TRANSMISSION_USERPASS" -a build/payload/temp%03d.torrent\n' % i
						payload += 'df\n'
						payload += 'transmission-remote "$TRANSMISSION_SERVER" -n "$TRANSMISSION_USERPASS" -l\n'
						payload += 'ls -l -R -L "$TRANSMISSION_DIRECTORY"\n'
						zf.writestr("payload.sh", payload)
					else:
						usage()
					# Mark the files as having been created on Windows so that
					# Unix permissions are not inferred as 0000
					for zfile in zf.filelist:
						zfile.create_system = 0
					zf.close()
				fp.seek(0)
				payload = fp.read()
				ext = 'zip'
	else:
		print("* Operation: upload file %s and run" % sys.argv[1])
		payload = open(sys.argv[1], 'rb').read()
		ext = sys.argv[1]
		i = ext.rfind(".")
		ext = ext[i+1:] if i >= 0 else ""

	# load config files
	config = json.load(open("build/config.json", 'r'))
	keys = json.load(open("build/client_key.json", 'r'))

	payload = create_payload(payload, ext, keys)

	inputs = create_inputs(config)

	for input in inputs:
		print("Upload the payload to %s" % input.name)
		input.upload_payload(payload, "input_box", "Encrypted message")
		return 0

	print("WARNING: there are nowhere to upload payload")


if __name__ == '__main__':
	main()
